#!/usr/bin/env python3
"""
修复EmbedTrack数据结构的脚本
"""

import os
import shutil
from pathlib import Path

def clean_and_reset_data(data_path):
    """
    清理已处理的数据并重置为原始状态
    
    Args:
        data_path (str): 数据目录路径
    """
    data_path = Path(data_path)
    
    print(f"清理数据目录: {data_path}")
    
    # 备份当前状态（可选）
    backup_path = data_path.parent / f"{data_path.name}_backup"
    if backup_path.exists():
        shutil.rmtree(backup_path)
    
    # 删除已处理的train和val目录
    train_dirs = [d for d in data_path.iterdir() if d.name.startswith('train_')]
    val_dirs = [d for d in data_path.iterdir() if d.name.startswith('val_')]
    
    for train_dir in train_dirs:
        print(f"删除: {train_dir}")
        shutil.rmtree(train_dir)
    
    for val_dir in val_dirs:
        print(f"删除: {val_dir}")
        shutil.rmtree(val_dir)
    
    # 删除train和val主目录（如果存在）
    if (data_path / "train").exists():
        shutil.rmtree(data_path / "train")
    if (data_path / "val").exists():
        shutil.rmtree(data_path / "val")
    
    print("数据清理完成")

def check_original_data_structure(data_path):
    """
    检查原始数据结构是否符合CTC格式
    
    Args:
        data_path (str): 数据目录路径
    """
    data_path = Path(data_path)
    
    print(f"检查数据结构: {data_path}")
    
    # 查找数字命名的目录（图像序列）
    image_dirs = [d for d in data_path.iterdir() if d.is_dir() and d.name.isdigit()]
    
    # 查找_ST或_GT结尾的目录（注释）
    annotation_dirs = [d for d in data_path.iterdir() 
                      if d.is_dir() and (d.name.endswith('_ST') or d.name.endswith('_GT'))]
    
    print(f"找到图像序列目录: {[d.name for d in image_dirs]}")
    print(f"找到注释目录: {[d.name for d in annotation_dirs]}")
    
    # 检查每个目录的内容
    for img_dir in image_dirs:
        files = list(img_dir.glob("*.tif"))
        print(f"  {img_dir.name}: {len(files)} 个图像文件")
    
    for ann_dir in annotation_dirs:
        subdirs = [d for d in ann_dir.iterdir() if d.is_dir()]
        print(f"  {ann_dir.name}: 子目录 {[d.name for d in subdirs]}")
        for subdir in subdirs:
            files = list(subdir.glob("*.tif"))
            txt_files = list(subdir.glob("*.txt"))
            print(f"    {subdir.name}: {len(files)} 个掩码文件, {len(txt_files)} 个txt文件")
    
    return image_dirs, annotation_dirs

def fix_data_preparation():
    """
    修复数据准备过程的主函数
    """
    # 替换为您的实际数据路径
    data_base_path = Path("d:/CodeRode/project/EmbedTrack-master/data")
    dataset_name = "3t3"  # 根据您的实际数据集名称调整
    
    dataset_path = data_base_path / dataset_name
    
    if not dataset_path.exists():
        print(f"错误: 数据集路径不存在: {dataset_path}")
        return
    
    print("=== 步骤1: 清理已处理的数据 ===")
    clean_and_reset_data(dataset_path)
    
    print("\n=== 步骤2: 检查原始数据结构 ===")
    image_dirs, annotation_dirs = check_original_data_structure(dataset_path)
    
    if not image_dirs:
        print("警告: 没有找到图像序列目录（数字命名的目录）")
        print("请确保您的原始CTC数据集结构正确")
        return
    
    if not annotation_dirs:
        print("警告: 没有找到注释目录（_ST或_GT结尾的目录）")
        print("请确保您的原始CTC数据集包含注释数据")
        return
    
    print("\n=== 数据结构看起来正确，可以重新运行prepare_ctc_data ===")
    
    # 提供重新运行的代码示例
    print("\n建议的重新运行代码:")
    print(f"""
from embedtrack.datasets.prepare_data import prepare_ctc_data

# 重新准备数据
prepare_ctc_data(
    source_path="{dataset_path}",
    result_path="{data_base_path}",
    keep_st=True,  # 或False，根据您的需要
    val_split=0.1,
    sub_dir_names=["46"]  # 根据您实际的序列号调整
)
""")

if __name__ == "__main__":
    fix_data_preparation()